import tqdm
import numpy as np
from loaddata import LoadData

# 计算图像数据的均值和标准差
def compute_mean_std(data, data_path):
    data = LoadData(data, data_path)
    mean_r = 0
    mean_g = 0
    mean_b = 0
    
    for i in tqdm.tqdm(data):
        # i 对于kitti是两张图，一张img，一张mask，用i[0]取img
        # mean_r += np.mean(i[0][:, :, 0])
        # mean_g += np.mean(i[0][:, :, 1])
        # mean_b += np.mean(i[0][:, :, 2])

        # 如果处理其他数据只返回图像数据的话不需要
        mean_r += np.mean(i[:, :, 0])
        mean_g += np.mean(i[:, :, 1])
        mean_b += np.mean(i[:, :, 2])
        
    mean_r /= len(data)
    mean_g /= len(data)
    mean_b /= len(data)
    
    diff_r = 0
    diff_g = 0
    diff_b = 0
    N = 0
    
    for i in tqdm.tqdm(data):
        # 同上
        # diff_r += np.sum(np.power(i[0][:, :, 0] - mean_r, 2))
        # diff_g += np.sum(np.power(i[0][:, :, 1] - mean_g, 2))
        # diff_b += np.sum(np.power(i[0][:, :, 2] - mean_b, 2))
        # N += np.prod(i[0][:, :, 0].shape)

        diff_r += np.sum(np.power(i[:, :, 0] - mean_r, 2))
        diff_g += np.sum(np.power(i[:, :, 1] - mean_g, 2))
        diff_b += np.sum(np.power(i[:, :, 2] - mean_b, 2))
        N += np.prod(i[:, :, 0].shape)
        
    std_r = np.sqrt(diff_r / N)
    std_g = np.sqrt(diff_g / N)
    std_b = np.sqrt(diff_b / N)
    
    mean = [mean_r.item() / 255.0, mean_g.item() / 255.0, mean_b.item() / 255.0]
    std = [std_r.item() / 255.0, std_g.item() / 255.0, std_b.item() / 255.0]

    return mean, std

if __name__ == "__main__":
    data = 'paris'
    data_path = '/workspace/wzj/revisitop/data/datasets/rparis6k'
    mean, std = compute_mean_std(data, data_path)
    print(mean)
    print(std)